#include <cmath>
#include <sstream>

#include "IBMModel1.hh"

double prob_smooth = 1e-7;
double IBMModel1::smoothed_p_s_given_t(int s, int t){
  double p = -1.0;

  if((p_s_given_t.find(s) == p_s_given_t.end()) ||
     (p_s_given_t[s].find(t) == p_s_given_t[s].end())){
    p = prob_smooth;
  }else{
    p = max(p_s_given_t[s][t], prob_smooth);
  }

  return p;
}

double IBMModel1::viterbi_inner(SentencePair &pair, bool verbose){
  int J = pair.source_ids.size();
  int I = pair.target_ids.size();
  double total_score = 0.0;

  for(int j = 0; j < J; j++){
    int s = pair.source_ids[j];
    int max_i = 0;
    double max_log_p = log(smoothed_p_s_given_t(s, pair.target_ids[0]));
    for(int i = 1; i < I; i++){
      int t = pair.target_ids[i];
      double log_p = log(smoothed_p_s_given_t(s, t));
      if(log_p > max_log_p){
	max_i = i;
	max_log_p = log_p;
      }
    }
    pair.alignment.insert(std::pair<int, int>(j, max_i));
    total_score += max_log_p;
  }
  
  return exp(total_score);
}

void IBMModel1::update_em(Bitext &bitext, Model *p_model){
  // initialize to zero parameter counts
  map<int, map<int, double> > counts_J_I;
  map<int, double>            counts_I;
  map<int, map<int, double> > counts_s_t;
  map<int, double>            counts_t;

  // for each item in bitext
  bitext.rewind();
  SentencePair pair;
  ExpectedCounts counts;
  while(bitext.next(pair)){
    // get expected counts of p_model
    counts.reset(pair);
    p_model->get_expected_counts(bitext, pair, counts);

    // add to counts of parameters:
    // counts_J_I
    counts_J_I[pair.J][pair.I] += 1;

    // c_I
    counts_I[pair.I] += 1;

    for(int j = 0; j < pair.J; j++){
      int s = pair.source_ids[j];
      for(int i = 0; i < pair.I; i++){
	int t = pair.target_ids[i];
	double c = counts[j][i];
	// c_s_t
	counts_s_t[s][t] += c;
	// c_t
	counts_t[t] += c;
      }
    }
  }

  // recalculate parameters:
  // p_j_given_i
  for(map<int, map<int, double> >::iterator itr = counts_J_I.begin(); itr != counts_J_I.end(); itr++){
    int J = itr->first;
    for(map<int, double>::iterator itr2 = itr->second.begin(); itr2 != itr->second.end(); itr2++){
      int I = itr2->first;
      double count_J_I = itr2->second;
      p_J_given_I[J][I] = count_J_I / counts_I[I];
    }
  }

  static int iteration = 0;
  stringstream i_s;
  i_s << iteration++;
  string p_s_given_t_dump_filename = "p_s_given_t." + i_s.str();
  ofstream p_s_given_t_dump;
  p_s_given_t_dump.open(p_s_given_t_dump_filename.c_str());
  for(map<int, map<int, double> >::iterator itr = counts_s_t.begin(); itr != counts_s_t.end(); itr++){
    int s = itr->first;
    for(map<int, double>::iterator itr2 = itr->second.begin(); itr2 != itr->second.end(); itr2++){
      int t = itr2->first;
      // p_s_given_t
      if((counts_s_t[s][t] > 1e-10) &&
	 (counts_t[t] > 1e-10)){
	double p = counts_s_t[s][t] / counts_t[t];
	if(p < prob_smooth){
	  prob_smooth = p - 1e-10;
	}
	if(isnan(p)){
	  cout << "got a nan! counts_s_t[" << s << "][" << t << "] = " << counts_s_t[s][t] << ", counts_t[" << t << "] = " << counts_t[t] << endl;
	}
	p_s_given_t_dump << bitext.source_vocabulary.map_to_token(s) << " " << bitext.target_vocabulary.map_to_token(t) << " " << p << endl;
	p_s_given_t[s][t] = p;
      }
    }
  }
  p_s_given_t_dump.close();
}

void IBMModel1::transfer(Bitext &bitext, Model *p_model){
  // update as in E-M
  update_em(bitext, p_model);
}

void IBMModel1::iterate(Bitext &bitext){
  update_em(bitext, this);
}

void IBMModel1::get_raw_expected_counts(Bitext &bitext, SentencePair &pair, ExpectedCounts &counts){
  map<int, map<int, double> > alpha;
  map<int, map<int, double> > beta;

  if((pair.J > 0) && (pair.I > 0)){ 
    for(int j = 0; j < pair.J; j++){
      int s = pair.source_ids[j];
      double sum = 0.0;
      for(int i = 0; i < pair.I; i++){
	int t = pair.target_ids[i];
	sum += p_s_given_t[s][t];
      }
      for(int i = 0; i < pair.I; i++){
	int t = pair.target_ids[i];
	counts[j][i] = pair.n * (p_s_given_t[s][t] / sum);
      }
    }
  }
}
